#!/bin/env python

import http.client as http_client
from urllib.parse import urlparse, urlencode

import json
import ssl
import sys
import zlib


def _ssl_context():
    ctx = ssl.create_default_context()
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE
    return ctx


def req(url):
    conn = http_client.HTTPSConnection(
        'marketplace.visualstudio.com', 443, _ssl_context())


def get_vsix_versions(vsix_id):
    headers = {
        "Accept": "application/json;api-version=7.2-preview.1;excludeUrls=true",
        "Content-Type": "application/json",
        "Referer": f"https://marketplace.visualstudio.com/items?itemName={vsix_id}",
        "Referrer-Policy": "strict-origin-when-cross-origin"
    }

    body = f"""
{{
    "assetTypes":null,
    "filters":[
        {{
            "criteria": [
                {{"filterType": 7, "value":"{vsix_id}"}}
            ],
            "direction":2,
            "pageSize":100,
            "pageNumber":1,
            "sortBy":0,
            "sortOrder":0,
            "pagingToken":null
        }}
    ],
    "flags":2151
}}"""

    conn = http_client.HTTPSConnection('marketplace.visualstudio.com')
    conn.request('POST', '/_apis/public/gallery/extensionquery',
                 body, headers=headers)

    resp = conn.getresponse()
    if resp.status != 200:
        raise RuntimeError('Error: %s' % resp.reason)

    content_encoding = resp.getheader('Content-Encoding', '').lower()
    data = resp.read()
    if content_encoding in ('gzip', 'deflate'):
        if content_encoding == 'gzip':
            decompressor = zlib.decompressobj(16 + zlib.MAX_WBITS)
            data = decompressor.decompress(data)
        elif content_encoding == 'deflate':
            decompressor = zlib.decompressobj()
            data = decompressor.decompress(data)
    #print(data.decode('utf-8'))

    json_data = json.loads(data)
    print(json.dumps(json_data, indent=4, ensure_ascii=False))

    versions = json_data['results'][0]['extensions'][0]["versions"]

    urls = []
    for v in versions:
        for f in v["files"]:
            if f["assetType"] == "Microsoft.VisualStudio.Services.VSIXPackage":
                platform = v.get("targetPlatform", "")
                url = f.get("source", "")
                print(f"{v['version']} {platform} {url}")
                urls.append((v['version'], platform, url))
    return urls

def get_vsix_url(vsix_id, version=None, platforms=[]):
    infos = []

    urls = get_vsix_versions(vsix_id)

    if urls:
        if not version: # latest
            version = urls[0][0]

        for v, p, url in urls:
            if v != version:
                continue
            if platforms and p not in platforms:
                continue
            infos.append((v, p, url))
    
    return infos

def download(url, outfile):
    headers = {
        'Accept-Encoding': 'gzip, deflate',
    }

    p = urlparse(url)

    conn = http_client.HTTPSConnection(p.hostname)
    conn.request('GET', p.path, headers=headers)

    resp = conn.getresponse()
    if resp.status != 200:
        raise RuntimeError('Error: %s' % resp.reason)

    total_size = int(resp.getheader('Content-Length', 0))
    content_encoding = resp.getheader('Content-Encoding', '').lower()
    transfer_encoding = resp.getheader('Transfer-Encoding', '').lower()

    down_size = 0

    def show_process(outfile, down_size, total_size):
        if total_size:
            print(f'\r{outfile}: {down_size/total_size:.1%}', end='')
        else:
            print('\r{outfile}: {down_size} Bytes', end='')

    with open(outfile, 'wb') as f:
        decompressor = None
        if content_encoding in ('gzip', 'deflate'):
            if content_encoding == 'gzip':
                decompressor = zlib.decompressobj(16 + zlib.MAX_WBITS)
            elif content_encoding == 'deflate':
                decompressor = zlib.decompressobj()

        chunk_size = 8192

        while True:
            if transfer_encoding == 'chunked':
                chunk_size_str = resp.readline().strip()
                chunk_size = int(chunk_size_str, 16)
                if chunk_size == 0:
                    break

            chunk = resp.read(chunk_size)
            if not chunk:
                break

            down_size += len(chunk)

            if decompressor:
                chunk = decompressor.decompress(chunk)

            f.write(chunk)
            show_process(outfile, down_size, total_size)

        if decompressor:
            f.write(decompressor.flush())

        print()


def main():
    if len(sys.argv) < 2:
        print('Usage: %s <vsix_id> [version] [platform]' % sys.argv[0])
        return 1

    vsix_id = sys.argv[1]
    version = ''
    platform = ''

    if len(sys.argv) > 2:
        version = sys.argv[2]
    if len(sys.argv) > 3:
        platform = sys.argv[3]

    infos = get_vsix_url(vsix_id, version, platform.split(','))
    if infos:
        for info in infos:
            ver, plat, url = info
            if plat:
                outfile = f'{vsix_id}-{ver}@{plat}.vsix'
            else:
                outfile = f'{vsix_id}-{ver}.vsix'
            download(url, outfile)
    else:
        print(f'Error: f{vsix_id} {version} {platform} not found')


if __name__ == '__main__':
    main()
